#include "gc.h"

// graph.create() must be called before add terminals
void addTerminals (Graph3D& graph, std::vector<int> terminalComponentNoSet, 
		   int terminalLayerNo, VerboseLevel verbose = VERBOSE_LOW)
{
  for (GraphLayerNodeSet::iterator nitr = graph.nodes[terminalLayerNo].begin();
       nitr != graph.nodes[terminalLayerNo].end(); nitr++) {
    // Node 0 is not considered in raw graph
    if (nitr->status == NODE_NORMAL && nitr->componentNo > 0) {
      int rawNodeNo = 
	graph.rawNodeCorrespondence[terminalLayerNo - graph.bodyStartLayerNo][nitr->no];
      bool isSrc = false;
      for (std::vector<int>::iterator citr = terminalComponentNoSet.begin(); 
	   citr != terminalComponentNoSet.end(); citr++) {
	if (nitr->componentNo == *citr) {
	  graph.body->add_tweights(rawNodeNo, FLOAT_INF, 0.0);
	  isSrc = true;
	  if (verbose == VERBOSE_HIGH) {
	    std::cout << "Add terminals: SOURCE <--> node " << nitr->no 
		      << " @ layer " << terminalLayerNo << std::endl;
	  }
	  break;
	}
      }
      if (!isSrc) {
	graph.body->add_tweights(rawNodeNo, 0.0, FLOAT_INF);
	if (verbose == VERBOSE_HIGH) {
	  std::cout << "Add terminals: SINK <--> node " << nitr->no 
		    << " @ layer " << terminalLayerNo << std::endl;
	}
      }
    }
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
double cut (Graph3D& graph, SegmentationGraph& segGraph, SegmentationType segType, 
	  VerboseLevel verbose = VERBOSE_LOW)
{
  double flow = graph.body->maxflow();
  int rawNodeNum = graph.nodeCorrespondence.size();
  for (int rawNodeNo = 0; rawNodeNo < rawNodeNum; rawNodeNo++) {
    if (graph.body->what_segment(rawNodeNo, RawGraph::SINK) 
	== RawGraph::SOURCE) {
      SegmentationNode segNode(graph.nodeCorrespondence[rawNodeNo].layerNo, 
			       graph.nodeCorrespondence[rawNodeNo].nodeNo, 
			       segType);
      segGraph.addSegNode(segNode);
    }
  }
  if (verbose >= VERBOSE_MEDIUM) {
    std::cout << "Graphcuts: max flow = " << flow << std::endl;
    std::cout << "Graphcuts: " << segGraph.segNodeSet.size() 
	      << " SOURCE nodes" << std::endl;
    std::cout << "Graphcuts: " << rawNodeNum - segGraph.segNodeSet.size()
	      << " SINK nodes" << std::endl;
  }
  return flow;
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void segment (Graph3D& graph, std::vector<int> fwdStartComponentNoSet, 
	      int startLayerNo, int layerNum, int pileLayerNum, 
	      SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	      double& fwdMaxFlow, double& bwdMaxFlow, 
	      VerboseLevel verbose)
{
  std::vector<int> nextStartComponentNoSet = fwdStartComponentNoSet;
  int endLayerNo = startLayerNo + layerNum - 1;
  int pileStartLayerNo;
  int pileEndLayerNo;
  // Forward cut
  pileStartLayerNo = startLayerNo;
  pileEndLayerNo = pileStartLayerNo + pileLayerNum - 1 > endLayerNo? 
    endLayerNo: pileStartLayerNo + pileLayerNum - 1;
  while (pileStartLayerNo < pileEndLayerNo) {
    int pileCurrentLayerNum = pileEndLayerNo - pileStartLayerNo + 1;
    // Create graph and cut
    graph.create(pileStartLayerNo, pileCurrentLayerNum, GRAPH_FWD);
    addTerminals(graph, nextStartComponentNoSet, pileStartLayerNo, verbose);
    fwdMaxFlow = cut(graph, fwdSegGraph, SEG_FWD, verbose);
    nextStartComponentNoSet.clear();
    // Find start components for next
    // Assume all elements in fwdSegGraph are in layer-increasing order
    int cutLastLayerNo = fwdSegGraph.segNodeSet.back().layerNo;
    for (SegGraphSegNodeSet::const_reverse_iterator itr = 
	   fwdSegGraph.segNodeSet.rbegin(); 
	 itr != fwdSegGraph.segNodeSet.rend(); itr++) {
      if (itr->layerNo == cutLastLayerNo) {
	nextStartComponentNoSet.push_back(graph.nodes[itr->layerNo][itr->nodeNo].componentNo);
      }
      else {
	break;
      }
    }
    if (cutLastLayerNo < pileEndLayerNo) {
      pileStartLayerNo = cutLastLayerNo;
      break;
    }
    pileStartLayerNo = pileEndLayerNo;
    pileEndLayerNo = pileStartLayerNo + pileLayerNum - 1 > endLayerNo? 
      endLayerNo: pileStartLayerNo + pileLayerNum - 1;
  }
  // Backward cut
  pileEndLayerNo = pileStartLayerNo - pileLayerNum + 1 < startLayerNo? 
    startLayerNo: pileStartLayerNo - pileLayerNum + 1;
  while (pileStartLayerNo > pileEndLayerNo) {
    int pileCurrentLayerNum = pileStartLayerNo - pileEndLayerNo + 1;
    // Create graph and cut
    graph.create(pileStartLayerNo, pileCurrentLayerNum, GRAPH_BWD);
    addTerminals(graph, nextStartComponentNoSet, pileStartLayerNo, verbose);
    bwdMaxFlow = cut(graph, bwdSegGraph, SEG_BWD, verbose);
    nextStartComponentNoSet.clear();
    // Find start components for next
    // Assume all elements in bwdSegGraph are in layer-decreasing order
    int cutLastLayerNo = bwdSegGraph.segNodeSet.back().layerNo;
    for (SegGraphSegNodeSet::const_reverse_iterator itr = 
	   bwdSegGraph.segNodeSet.rbegin(); 
	 itr != bwdSegGraph.segNodeSet.rend(); itr++) {
      if (itr->layerNo == cutLastLayerNo) {
	nextStartComponentNoSet.push_back(graph.nodes[itr->layerNo][itr->nodeNo].componentNo);
      }
      else {
	break;
      }
    }
    if (cutLastLayerNo > pileEndLayerNo) {
      break;
    }
    pileStartLayerNo = pileEndLayerNo;
    pileEndLayerNo = pileStartLayerNo - pileLayerNum + 1 < startLayerNo? 
      startLayerNo: pileStartLayerNo - pileLayerNum + 1;
  }
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
void segment (Graph3D& graph, std::vector<int> fwdStartComponentNoSet, 
	      int startLayerNo, int layerNum, 
	      SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	      double& fwdMaxFlow, double& bwdMaxFlow, VerboseLevel verbose)
{
  graph.create(startLayerNo, layerNum);
  addTerminals(graph, fwdStartComponentNoSet, startLayerNo, verbose);
  fwdMaxFlow = cut(graph, fwdSegGraph, SEG_FWD, verbose);
  std::vector<int> bwdStartComponentNoSet;
  int lastLayerNo = -1;
  for (SegGraphSegNodeSet::const_iterator itr = fwdSegGraph.segNodeSet.begin(); 
       itr != fwdSegGraph.segNodeSet.end(); itr++) {
    if (itr->layerNo > lastLayerNo) {
      lastLayerNo = itr->layerNo;
    }
  }
  for (SegGraphSegNodeSet::const_iterator itr = fwdSegGraph.segNodeSet.begin(); 
       itr != fwdSegGraph.segNodeSet.end(); itr++) {
    if (itr->layerNo == lastLayerNo) {
      bwdStartComponentNoSet.push_back(graph.nodes[itr->layerNo][itr->nodeNo].componentNo);
    }
  }
  graph.create(startLayerNo, layerNum);
  addTerminals(graph, bwdStartComponentNoSet, lastLayerNo, verbose);
  bwdMaxFlow = cut(graph, bwdSegGraph, SEG_BWD, verbose);
}
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
SegGraphSegNodeSet::iterator locateSegNode (SegmentationGraph& segGraph, 
					    SegmentationNode const& segNode)
{
  SegGraphSegNodeSet::iterator itr;
  for (itr = segGraph.segNodeSet.begin(); itr != segGraph.segNodeSet.end(); 
       itr++) {
    if (itr->isEqual(segNode)) {
      break;
    }
  }
  return itr;
}


inline SegmentationType determineSegType (SegmentationType type0, 
					  SegmentationType type1)
{
  if (type0 == SEG_NONE) {
    return type1;
  }
  else if (type1 == SEG_NONE) {
    return type0;
  }
  else if (type0 == SEG_BOTH || type1 == SEG_BOTH) {
    return SEG_BOTH;
  }
  else if (type0 == type1) {
    return type0;
  }
  else {
    return SEG_BOTH;
  }
}



void extractGroup (Volume const& volume, Graph3D const& graph, int startLayerNo, 
		   SegmentationGraph fwdSegGraph, SegmentationGraph bwdSegGraph, 
		   std::vector<int> fwdStartComponentNoSet, 
		   SegmentationGraph& resSegGraph)
{
  std::queue<SegmentationNode> gQueue;
  for (std::vector<int>::const_iterator citr = fwdStartComponentNoSet.begin(); 
       citr != fwdStartComponentNoSet.end(); citr++) {
    SegmentationNode csn(startLayerNo, *citr);
    SegGraphSegNodeSet::iterator fwd_csn_itr = locateSegNode(fwdSegGraph, csn);
    SegmentationType fwdSegType = 
      fwd_csn_itr == fwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_FWD;
    SegGraphSegNodeSet::iterator bwd_csn_itr = locateSegNode(bwdSegGraph, csn);
    SegmentationType bwdSegType = 
      bwd_csn_itr == bwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_BWD;
    csn.segType = determineSegType(fwdSegType, bwdSegType);
    if (csn.segType != SEG_NONE) {
      gQueue.push(csn);
      if (fwdSegType == SEG_FWD) {
	fwdSegGraph.segNodeSet.erase(fwd_csn_itr);
      }
      if (bwdSegType == SEG_BWD) {
	bwdSegGraph.segNodeSet.erase(bwd_csn_itr);
      }      
    }
  }
  while (!gQueue.empty()) {
    SegmentationNode sn = gQueue.front();
    gQueue.pop();
    resSegGraph.segNodeSet.push_back(sn);
    int layerNo = sn.layerNo;
    int nodeNo = sn.nodeNo;
    // Add forward edges end node
    // Adjacent edges
    for (NodeEdgeSet::const_iterator itr = 
	   graph.nodes[layerNo][nodeNo].edges.begin(); 
	 itr != graph.nodes[layerNo][nodeNo].edges.end(); 
	 itr++) {
      SegmentationNode csn(layerNo + 1, itr->endNodeNo);
      SegGraphSegNodeSet::iterator fwd_csn_itr = locateSegNode(fwdSegGraph, csn);
      SegmentationType fwdSegType = 
	fwd_csn_itr == fwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_FWD;
      SegGraphSegNodeSet::iterator bwd_csn_itr = locateSegNode(bwdSegGraph, csn);
      SegmentationType bwdSegType = 
	bwd_csn_itr == bwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_BWD;
      csn.segType = determineSegType(fwdSegType, bwdSegType);
      if (csn.segType != SEG_NONE) {
	gQueue.push(csn);
	if (fwdSegType == SEG_FWD) {
	  fwdSegGraph.segNodeSet.erase(fwd_csn_itr);
	}
	if (bwdSegType == SEG_BWD) {
	  bwdSegGraph.segNodeSet.erase(bwd_csn_itr);
	}      
      }
    }
    // Skip edges
    for (NodeEdgeSet::const_iterator itr = 
	   graph.nodes[layerNo][nodeNo].skipEdges.begin(); 
	 itr != graph.nodes[layerNo][nodeNo].skipEdges.end(); 
	 itr++) {
      SegmentationNode csn(layerNo + 2, itr->endNodeNo);
      SegGraphSegNodeSet::iterator fwd_csn_itr = locateSegNode(fwdSegGraph, csn);
      SegmentationType fwdSegType = 
	fwd_csn_itr == fwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_FWD;
      SegGraphSegNodeSet::iterator bwd_csn_itr = locateSegNode(bwdSegGraph, csn);
      SegmentationType bwdSegType = 
	bwd_csn_itr == bwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_BWD;
      csn.segType = determineSegType(fwdSegType, bwdSegType);
      if (csn.segType != SEG_NONE) {
	gQueue.push(csn);
	if (fwdSegType == SEG_FWD) {
	  fwdSegGraph.segNodeSet.erase(fwd_csn_itr);
	}
	if (bwdSegType == SEG_BWD) {
	  bwdSegGraph.segNodeSet.erase(bwd_csn_itr);
	}      
      }
    }
    // Add reverse edges end node
    // Adajacent edges
    for (NodeEdgeSet::const_iterator itr = 
	   graph.nodes[layerNo][nodeNo].reverseEdges.begin(); 
	 itr != graph.nodes[layerNo][nodeNo].reverseEdges.end(); 
	 itr++) {
      SegmentationNode csn(layerNo - 1, itr->endNodeNo);
      SegGraphSegNodeSet::iterator fwd_csn_itr = locateSegNode(fwdSegGraph, csn);
      SegmentationType fwdSegType = 
	fwd_csn_itr == fwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_FWD;
      SegGraphSegNodeSet::iterator bwd_csn_itr = locateSegNode(bwdSegGraph, csn);
      SegmentationType bwdSegType = 
	bwd_csn_itr == bwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_BWD;
      csn.segType = determineSegType(fwdSegType, bwdSegType);
      if (csn.segType != SEG_NONE) {
	gQueue.push(csn);
	if (fwdSegType == SEG_FWD) {
	  fwdSegGraph.segNodeSet.erase(fwd_csn_itr);
	}
	if (bwdSegType == SEG_BWD) {
	  bwdSegGraph.segNodeSet.erase(bwd_csn_itr);
	}      
      }
    }
    // Skip edges
    for (NodeEdgeSet::const_iterator itr = 
	   graph.nodes[layerNo][nodeNo].reverseSkipEdges.begin(); 
	 itr != graph.nodes[layerNo][nodeNo].reverseSkipEdges.end(); 
	 itr++) {
      SegmentationNode csn(layerNo - 2, itr->endNodeNo);
      SegGraphSegNodeSet::iterator fwd_csn_itr = locateSegNode(fwdSegGraph, csn);
      SegmentationType fwdSegType = 
	fwd_csn_itr == fwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_FWD;
      SegGraphSegNodeSet::iterator bwd_csn_itr = locateSegNode(bwdSegGraph, csn);
      SegmentationType bwdSegType = 
	bwd_csn_itr == bwdSegGraph.segNodeSet.end()? SEG_NONE: SEG_BWD;
      csn.segType = determineSegType(fwdSegType, bwdSegType);
      if (csn.segType != SEG_NONE) {
	gQueue.push(csn);
	if (fwdSegType == SEG_FWD) {
	  fwdSegGraph.segNodeSet.erase(fwd_csn_itr);
	}
	if (bwdSegType == SEG_BWD) {
	  bwdSegGraph.segNodeSet.erase(bwd_csn_itr);
	}      
      }
    }
  }
  resSegGraph.sort();
}
